Skip to content

fix(algorithms): raise instead of asserting on the skipped-reference-KL pairing - #3854

Open
tianyi-zhang-02 wants to merge 5 commits into
NVIDIA-NeMo:mainfrom
tianyi-zhang-02:fix-reference-kl-assert
Open

fix(algorithms): raise instead of asserting on the skipped-reference-KL pairing#3854
tianyi-zhang-02 wants to merge 5 commits into
NVIDIA-NeMo:mainfrom
tianyi-zhang-02:fix-reference-kl-assert

Conversation

@tianyi-zhang-02

@tianyi-zhang-02 tianyi-zhang-02 commented Aug 26, 2026

Copy link
Copy Markdown
Contributor

What does this PR do?

Replaces skipped-reference-logprob configuration asserts in synchronous PPO and GRPO with ValueErrors, so the safety check still runs under python -O.

The GRPO validation also runs at the start of setup, before logger, dataloader, checkpoint, or cluster initialization. The train-loop check remains as a defensive guard.

Validation

Final SHA: d64af0e6d7dc4c29c24fe35df4ad2ae1a20d28c6, based on upstream main at ccbcd4cc5.

  • tests/unit/algorithms/test_reference_kl_guard.py: 5 passed
  • setup coverage asserts the user-facing error occurs before cluster-side work
  • Ruff check and format-check passed

This is setup/config validation, so no GPU workload is required.

…KL pairing

Sync PPO and sync GRPO both guard skip_reference_policy_logprobs_calculation
against a non-zero reference_policy_kl_penalty with a bare, message-less

    assert master_config.loss_fn.reference_policy_kl_penalty == 0

python -O removes it entirely, so the run proceeds to train against a KL term
whose reference logprobs were never computed. Even with asserts enabled, a
message-less one says nothing about which two settings conflict.

async_ppo_train already does this as an if + ValueError. This is the same
defect class yuki-97 raised on NVIDIA-NeMo#3262 ("assert backend == 'vllm' gets stripped
under python -O"), which was converted there but left in these two siblings.

Both now use the async form, verbatim, so the three read alike.

Tests run under python -O as well as normally, which is what separates this
from a cosmetic change: with the assert restored, both modes fail.

Signed-off-by: Tianyi Zhang <zhangtianyi975@gmail.com>
Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
@tianyi-zhang-02
tianyi-zhang-02 requested review from a team as code owners August 26, 2026 18:49
@copy-pr-bot

copy-pr-bot Bot commented Aug 26, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

…ed one

grpo_train_sync's assert is unreachable in practice: grpo.setup checks the
same pairing before the train loop starts, so on GRPO that is the one a user
hits. It was an assert too, so under python -O both were stripped and nothing
was left -- converting only the train-loop copy would have fixed the copy
nobody reaches.

PPO is the other way round: there is no setup-time equivalent, so the one in
ppo_train is the only guard and converting it was already right.

Keeps the grpo_sync conversion as well. It costs nothing and the three sites
now read alike, which was the point.

Test asserts on grpo.setup's source rather than calling it, since calling it
needs a cluster. What it pins is that the reachable guard is not an assert.

Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
@tianyi-zhang-02
tianyi-zhang-02 force-pushed the fix-reference-kl-assert branch from a3488d8 to 4bf8b1d Compare August 28, 2026 15:28
@svcnvidia-nemo-ci svcnvidia-nemo-ci added the waiting-on-maintainers Waiting on maintainers to respond label Aug 28, 2026
Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

community-request waiting-on-maintainers Waiting on maintainers to respond

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants